#include "wii_mmgr.h"

#ifdef _PLATFORM_WII
#include <revolution.h>
#include <revolution/mem.h>
#include <cassert>
#include <os.h>
#include <Stdio.h>
#include <pool_alloc.h>
#include <vector>
#include <stdlib.h> 


#include "common_macros.h"


//#if __ide_target("Debug")
#define		ENABLE_MEMORY_DEBUG 1
//#endif // debug target

#ifdef ENABLE_MEMORY_DEBUG
	struct		StringNode;
	struct		MemoryTrack
	{
		unsigned int	Size;
		StringNode*		pFile;
		StringNode*		pFunc;
		unsigned int	Line;
		// treat it as a two-way linked list so we can easily add/delete nodes
		MemoryTrack*	pPrev;
		MemoryTrack*	pNext;
	};
	static const int		PADDING = OSRoundUp32B(sizeof(MemoryTrack));
	static	const	char	*TrackFileName         = "unknown";
	static	const	char	*TrackFuncName 			= "unknown";
	static	unsigned int	TraceFuncLine             = 0;
	static MemoryTrack*		pBaseTrack = NULL;
#else
	static const int		PADDING = 0;
#endif

static bool				MMGRInit = false;
struct				MMGRHeapStats
{

	int				MaxMemory;
	int				Allocations;
	int				DeAllocations;
	int				MaxAllocations;
	int				CurrentUsage;
	int				HighestUsage;
	MEMHeapHandle	Handle;
	void*			pFrom;
	void*			pTo;
	// used for unit heaps only
	int				UnitC, UnitSize;
	
	int				GetAllocations()
	{
		return Allocations-DeAllocations;
	}
};

static MMGRHeapStats	HeapInfoA[HEAP_C];
static int				MMGR_HeapMisses;

/////////////////////////////////////////////////////////////
// We're wanting to track the full data of the allocation (file/line/func)
// But doing it in such a way as to remove any allocations from outside this file
struct	StringNode
{
	char*			Text;
	unsigned int	Key;
	StringNode*	pNext;
	void	 Init(const char* Input) 
	{
		pNext = NULL;
		Text = NULL;

		int		Size = strlen(Input);
		// don't let it go through our regular allocation engine
		Text = (char*)MEMAllocFromExpHeapEx(HeapInfoA[HEAP_MEM2].Handle, Size+1, DEFAULT_ALIGNMENT	);
		memcpy(Text, Input, Size+1);

		Key = GetKey(Text);
	};
	void			Delete()
	{
		if (pNext)	
		{
			pNext->Delete();
			MEMFreeToExpHeap(HeapInfoA[HEAP_MEM2].Handle, pNext);
		}
		if (Text)	MEMFreeToExpHeap(HeapInfoA[HEAP_MEM2].Handle, Text);
	}
	static unsigned int GetKey(const char* Input)
	{
		unsigned int	keyValue = 17;
		int				Size = strlen(Input);
		for (unsigned int i=0; i<Size; i++)
		{
			const char* pToken = &Input[i];
			keyValue *= 65599;
			keyValue = (keyValue ^ *pToken);
		}
		return( keyValue );
	}
};

static const int		LOOKUP_TABLE_SIZE = 1001;
static StringNode*		LookupMapA[LOOKUP_TABLE_SIZE] = {};

static void				LookupMapInit()
{
	for (unsigned int i=0; i<LOOKUP_TABLE_SIZE; i++) LookupMapA[i] = 0;
}

static StringNode*		LookupMapFindStr(const char* String)
{
	unsigned int Key = StringNode::GetKey(String);
	unsigned int Bin = Key % LOOKUP_TABLE_SIZE;
	StringNode*	pCurrent = LookupMapA[Bin];
	while (pCurrent)
	{
		if (pCurrent->Key == Key) return pCurrent;
		pCurrent = pCurrent->pNext;
	}
	return NULL;
}

static StringNode*		LookupMapAddString(const char* String)
{
	StringNode*	pNode 	= LookupMapFindStr(String);
	if (pNode) 
		return pNode;
	unsigned int Key 	= StringNode::GetKey(String);
	unsigned int Bin 	= Key % LOOKUP_TABLE_SIZE;
	pNode 				= (StringNode*)MEMAllocFromExpHeapEx(HeapInfoA[HEAP_MEM2].Handle, sizeof(StringNode), DEFAULT_ALIGNMENT	);
	pNode->Init(String);
	pNode->pNext		= LookupMapA[Bin];
	LookupMapA[Bin] 	= pNode;
	return pNode;
}




/////////////////////////////////
void*			CreateUnitHeap(MMGRHeapStats& HeapID, void* pFrom, int Count, unsigned int Size)
{
	HeapID.UnitC = Count;
	HeapID.UnitSize = Size;
	HeapID.MaxMemory = OSRoundUp32B(Size * Count);
	HeapID.Handle = MEMCreateUnitHeapEx(pFrom,  HeapID.MaxMemory, Size, DEFAULT_ALIGNMENT, MEM_HEAP_OPT_0_CLEAR);
	HeapID.pFrom = pFrom;
	HeapID.pTo =  (void*)OSRoundUp32B((unsigned char*)(pFrom) + HeapID.MaxMemory);
	return HeapID.pTo;
}

void*			CreateHeap(MMGRHeapStats& HeapID, void* pFrom, unsigned int Size)
{
	HeapID.MaxMemory = OSRoundUp32B(Size);
	HeapID.Handle = MEMCreateExpHeapEx(pFrom, Size, MEM_HEAP_OPT_0_CLEAR);
	HeapID.pFrom = pFrom;
	HeapID.pTo =  (void*)OSRoundUp32B((unsigned char*)(pFrom) + HeapID.MaxMemory);
	return HeapID.pTo;

}

MMGR_HEAPS	GetHeap(void* pAddr)
{
	for (int i=0; i<HEAP_C; i++)
	{
		if (pAddr >= HeapInfoA[i].pFrom && pAddr < HeapInfoA[i].pTo) return (MMGR_HEAPS)i;
	}
	return 	HEAP_UNKNOWN;
}

////////////////
// Direct access for C functions
void		WiiMMGR_Init()
{
	if (MMGRInit)
	{
		return;
	}
	memset(HeapInfoA, 0, sizeof(HeapInfoA));
	LookupMapInit();

  	OSInit();
	// initialize the basic memory manager
	void *lo1 = (void*)OSRoundUp32B(OSGetMEM1ArenaLo());
	void *hi1 = (void*)OSRoundUp32B(OSGetMEM1ArenaHi());	
	
	// Heap in MEM2
	void *lo2 = (void*)OSRoundUp32B(OSGetMEM2ArenaLo());
	void *hi2 = (void*)OSRoundUp32B(OSGetMEM2ArenaHi());	

	unsigned int Mem1 = (unsigned int)hi1-(unsigned int)lo1;
	unsigned int Mem2 = (unsigned int)hi2-(unsigned int)lo2;
	
	// create the unit heaps
	void* pCurrent = lo1;
	pCurrent = CreateUnitHeap(HeapInfoA[HEAP_UNIT_16], pCurrent, 22000, 16);
	pCurrent = CreateUnitHeap(HeapInfoA[HEAP_UNIT_32], pCurrent, 50000, 32);
	pCurrent = CreateUnitHeap(HeapInfoA[HEAP_UNIT_64], pCurrent, 25000, 64);
	pCurrent = CreateUnitHeap(HeapInfoA[HEAP_UNIT_256], pCurrent, 30000, 256);
	
	
	// create the main block of memory in MEM1 with whatever's left
	pCurrent = CreateHeap(HeapInfoA[HEAP_MEM1], pCurrent, (unsigned int)(hi1)-(unsigned int)(pCurrent));	
	
	// then create a single chunk that takes the whole of MEM2
	// Q: fake this to the 'correct' size of MEM2 on commercial sets, leave the rest as explicit debug?
	CreateHeap(HeapInfoA[HEAP_MEM2], lo2, (unsigned int)(hi2)-(unsigned int)(lo2));	

	// init the memory tracker
 	MMGR_HeapMisses = 0;
	MMGRInit = true;
}

void		WiiMMGR_Shutdown()
{
	if (!MMGRInit)
	{
		return;
	}
	for (int i=0; i<HEAP_C; i++)
	{
		if (HeapInfoA[i].Handle) 
		{
			if (MEMIsUnitHeap(HeapInfoA[i].Handle))
			{
				MEMDestroyUnitHeap( HeapInfoA[i].Handle);
			} else
			{
				MEMDestroyExpHeap(HeapInfoA[i].Handle);
			}	
		}
	}	
}

int			WiiMMGR_MemUsage(unsigned int& Mem1, unsigned int& Mem2)
{
	Mem1 = 0;
	Mem2 = 0;
	Mem1+=HeapInfoA[HEAP_MEM1].CurrentUsage;
	// unit heaps are all allocated from MEM1 as well
	Mem1+=HeapInfoA[HEAP_UNIT_16].CurrentUsage;
	Mem1+=HeapInfoA[HEAP_UNIT_32].CurrentUsage;
	Mem1+=HeapInfoA[HEAP_UNIT_64].CurrentUsage;
	Mem1+=HeapInfoA[HEAP_UNIT_256].CurrentUsage;
	// mem2 is just a single block	
	Mem2+=HeapInfoA[HEAP_MEM2].CurrentUsage;
	return 0;
}


unsigned int		WiiMMGR_SetOwner(const	char* file, unsigned int line,  const	char* func)
{
#ifdef ENABLE_MEMORY_DEBUG
	TrackFileName         	= file;
	TrackFuncName 		 	= func;
	TraceFuncLine           = line;
#else

#endif //ENABLE_MEMORY_DEBUG
	return 1;
}

#ifdef ENABLE_MEMORY_DEBUG
static int		SortBySize(const void* pA, const void* pB)
{
	const WiiAllocation* pNodeA = (const WiiAllocation*)pA;
	const WiiAllocation* pNodeB = (const WiiAllocation*)pB;
	return pNodeB->Size - pNodeA->Size;
}

void		WiiMMGR_PrintAll()
{
	char	Buffer[256];
	unsigned int	Count = 0;
	unsigned int	TotalMem1 = 0;
	unsigned int	TotalMem2 = 0;
	unsigned int	UnknownMem = 0;

	WiiAllocation*	AllocationA = (WiiAllocation*)WiiMMGR_ScratchPad();
	MemoryTrack* pNode = pBaseTrack;
	while (pNode)
	{
		char*	pFile = "";
		char*	pFunc = "";
		if (!pNode->pFile || !pNode->pFunc)
		{
		}
		else
		{
			pFile = pNode->pFile->Text;
			pFunc = pNode->pFunc->Text;
		}

		MMGR_HEAPS	src = GetHeap(pNode);
		if (src == HEAP_MEM1)
			TotalMem1 += pNode->Size;
		else if (src == HEAP_MEM2)
			TotalMem2 += pNode->Size;
		else
			UnknownMem += pNode->Size;


		//sprintf(Buffer, "%s:%i[%s] : %ikb\n", pFile, pNode->Line, pFunc, pNode->Size/1024);
		//OSReport(Buffer);

		// fill up the buffer
		AllocationA[Count].File = pFile;
		AllocationA[Count].Func = pFunc;
		AllocationA[Count].Line = pNode->Line;
		AllocationA[Count].Size = pNode->Size;

		Count++;
		pNode = pNode->pNext;
	}
	
	
	// sort and process?
	unsigned int 	Huge = 0;
	qsort(AllocationA, Count, sizeof(WiiAllocation), SortBySize);
	for (unsigned int i=0; i<Count; i++)
	{

		WiiAllocation*	pNode = &AllocationA[i];
		sprintf(Buffer, "%s:%i[%s] : %ikb\n", pNode->File, pNode->Line, pNode->Func, pNode->Size/1024);
		OSReport(Buffer);
		Huge ++;
		if (pNode->Size < 32 * 1024)		break;
	}
	sprintf(Buffer, "%i allocations, totalling %iMB (%ikb, %ikb, %ikb)\n", Count, (TotalMem1+ TotalMem2+UnknownMem)/(1024*1024), TotalMem1/1024, TotalMem2/1024, UnknownMem/1024);
	OSReport(Buffer);
	sprintf(Buffer, "%i allocations > 32Kb, %i < below. \n", Huge, Count-Huge);
	OSReport(Buffer);
}
#endif

void*		WiiMMGR_Alloc(MMGR_HEAPS HeapID,unsigned int Size, int Alignment)
{
	if (!MMGRInit)
		WiiMMGR_Init();
	static int 	Allocations = 0;
	

	void*	pData = NULL;	
	Allocations++;	
	if (HeapID == HEAP_MEM2 && Size <= 256 && Alignment == DEFAULT_ALIGNMENT)	// we've just alligned ourselves along the 4-byte boundries, which is fine for most. Some objects do require 32B alignment, though
	{
		HeapID = HEAP_UNKNOWN;
		if (Size <= 16)  { HeapID = HEAP_UNIT_16; Size = 16; } else
		if (Size <= 32)  { HeapID = HEAP_UNIT_32; Size = 32; } else
		if (Size <= 64)  { HeapID = HEAP_UNIT_64; Size = 64; } else
		if (Size <= 256) { HeapID = HEAP_UNIT_256; Size = 256;}

		if (HeapID < HEAP_C)
		{
			pData = MEMAllocFromUnitHeap(HeapInfoA[HeapID].Handle);
		}
		else
		{
			HeapID = HEAP_MEM2;	// restore this one
		}
		if (!pData)
		{
			int j=0;// failed to get memory!
		}
	}
	if (!pData)
	{
		if (HeapID <= HEAP_MEM2)	// only do this for normal heaps!
		{
			// unit heaps didn't catch it, allocate normally	
			#ifdef ENABLE_MEMORY_DEBUG
				int	PaddedSize = Size+PADDING;
				pData = MEMAllocFromExpHeapEx(HeapInfoA[HeapID].Handle, PaddedSize, Alignment);
			#else
				pData = MEMAllocFromExpHeapEx(HeapInfoA[HeapID].Handle, Size, Alignment);		
			#endif
		}
		else
		{
			HeapID = HEAP_MEM2;	// we just blew through one of the unit-heaps here!
			pData = MEMAllocFromExpHeapEx(HeapInfoA[HeapID].Handle, Size+PADDING, Alignment);
		}
		#ifdef ENABLE_MEMORY_DEBUG
		if (pData)
		{
			MemoryTrack* pHeader = (MemoryTrack*)pData;
			pData = ((char*)pData)+PADDING;	// advance it here... 
			pHeader->Size = Size;
			pHeader->Line = TraceFuncLine;
			pHeader->pFile = LookupMapAddString(TrackFileName);
			pHeader->pFunc = LookupMapAddString(TrackFuncName);
			
			pHeader->pNext = pBaseTrack;
			pHeader->pPrev = NULL;	// top node never has anything above it
			if (pBaseTrack)
			{
				pBaseTrack->pPrev = pHeader;
			}

			// replace on top of set
			pBaseTrack = pHeader;

			// clear it back out again!
			TrackFileName         	= "unknown";
			TrackFuncName 		 	= "unknown";
			TraceFuncLine           = 0;
		}
		#endif
	}
	if (!pData)
	{
		assert("Failed to allocate memory!");
	
		// log out memory as well?
//		const char*	Memory = WiiMMGR_ToString();
//		OSReport("OUT OF MEMORY!");
//		OSReport(Memory);
		WiiMMGR_PrintAll();
		return NULL;
	}
	// update stats
	HeapInfoA[HeapID].Allocations++;
	HeapInfoA[HeapID].MaxAllocations = UTIL_HIGH(HeapInfoA[HeapID].MaxAllocations, HeapInfoA[HeapID].GetAllocations());
#ifdef ENABLE_MEMORY_DEBUG
	HeapInfoA[HeapID].CurrentUsage+=Size;
#endif 
	HeapInfoA[HeapID].HighestUsage=UTIL_HIGH(HeapInfoA[HeapID].CurrentUsage, HeapInfoA[HeapID].HighestUsage);
	return pData;
}



void			WiiMMGR_Free(void* pData)
{
	if (pData == 0) return;	// quick exit; no action needed
	if (!MMGRInit)	WiiMMGR_Init();			// how did that happen?
	
	MMGR_HEAPS	src = GetHeap(pData);
	if (src > HEAP_C)	// oops
	{
		OSReport("Invalid Heap!");
		MMGR_HeapMisses++;
		return;	
	}
	
	if (src <= HEAP_MEM2)	// it's a normal chunk
	{
	
		#ifdef ENABLE_MEMORY_DEBUG
			MemoryTrack* pHeader = (MemoryTrack*)(((char*)pData) - PADDING);
			// update the nodes
			if (pHeader->pPrev)
			{
				pHeader->pPrev->pNext = pHeader->pNext;
			}
			if (pHeader->pNext)
			{
				pHeader->pNext->pPrev = pHeader->pPrev;
			}	
			if (pHeader == pBaseTrack)
			{
				pBaseTrack = pHeader->pNext;
			}
			
			HeapInfoA[src].CurrentUsage-=pHeader->Size;
			MEMFreeToExpHeap(HeapInfoA[src].Handle, pHeader);	// must delete from this source!
		#else
			MEMFreeToExpHeap(HeapInfoA[src].Handle, pData);
		#endif // ENABLE_MEMORY_DEBUG
	}
	else				// it's a unit-heap
	{
		MEMFreeToUnitHeap(HeapInfoA[src].Handle, pData);
		#ifdef ENABLE_MEMORY_DEBUG
		switch (src)
		{
			case HEAP_UNIT_16:	HeapInfoA[src].CurrentUsage-=16; break;
			case HEAP_UNIT_32:	HeapInfoA[src].CurrentUsage-=32; break;
			case HEAP_UNIT_64:	HeapInfoA[src].CurrentUsage-=64; break;
			case HEAP_UNIT_256:	HeapInfoA[src].CurrentUsage-=256; break;
		}
		#endif //ENABLE_MEMORY_DEBUG
		
	}
	
	HeapInfoA[src].DeAllocations++;
}

void*		WiiMMGR_Realloc(void* pArea, unsigned int	OldSize, unsigned int NewSize, int Alignment)
{
	// cache and reset the line/src/func?
	if (!pArea) 
		return WiiMMGR_Alloc(HEAP_MEM2, NewSize, Alignment);

	MMGR_HEAPS HeapID = GetHeap(pArea);
	if (HeapID > HEAP_MEM2)		HeapID = HEAP_MEM2;	// the alloc can handle this guess itself
	void*	pNew = WiiMMGR_Alloc(HeapID, NewSize, Alignment);
	if (!pNew) 
	{
		WiiMMGR_Free(pArea);
		return NULL;
	}
	if (!pArea) return pNew;	// essentially the same
	int		CpySize = UTIL_LOW(NewSize, OldSize);
	if (CpySize > 0)
	{
		memcpy(pNew, pArea, CpySize);
	}
	WiiMMGR_Free(pArea);
	return pNew;
}

// provide a debug output for strings
const char*	WiiMMGR_ToString()
{
	static char 	Buffer[1024];
	const char*		Format = 	"Total %ikb/%ikb in %i allocations (%i misses)\n"
							 	"MEM1 (%ikb/%ikb in %i units)\n"
							 	"MEM2 (%ikb/%ikb in %i units)\n"
							 	"UNIT16  (%ikb/%ikb in %i units, %i Max)\n"
							 	"UNIT32  (%ikb/%ikb in %i units, %i Max)\n"
							 	"UNIT64  (%ikb/%ikb in %i units, %i Max)\n"
							 	"UNIT256 (%ikb/%ikb in %i units, %i Max)\n";

	
	int				TotalMemoryUsed = 0;
	int				TotalMemoryReserved = 0;
	int				TotalAllocations =0;
	for (int i=0; i<HEAP_C; i++)
	{
		TotalMemoryUsed += HeapInfoA[i].CurrentUsage;
		TotalMemoryReserved += HeapInfoA[i].MaxMemory;
		TotalAllocations +=HeapInfoA[i].Allocations - HeapInfoA[i].DeAllocations;
	}	

	sprintf(Buffer, Format,
			TotalMemoryUsed/1024, TotalMemoryReserved/1024, TotalAllocations, MMGR_HeapMisses,		// global data
			HeapInfoA[HEAP_MEM1].CurrentUsage/1024, HeapInfoA[HEAP_MEM1].MaxMemory/1024, HeapInfoA[HEAP_MEM1].Allocations - HeapInfoA[HEAP_MEM1].DeAllocations,
			HeapInfoA[HEAP_MEM2].CurrentUsage/1024, HeapInfoA[HEAP_MEM2].MaxMemory/1024, HeapInfoA[HEAP_MEM2].Allocations - HeapInfoA[HEAP_MEM2].DeAllocations,
			HeapInfoA[HEAP_UNIT_16].CurrentUsage/1024, HeapInfoA[HEAP_UNIT_16].MaxMemory/1024, HeapInfoA[HEAP_UNIT_16].GetAllocations(), HeapInfoA[HEAP_UNIT_16].MaxAllocations,
			HeapInfoA[HEAP_UNIT_32].CurrentUsage/1024, HeapInfoA[HEAP_UNIT_32].MaxMemory/1024, HeapInfoA[HEAP_UNIT_32].GetAllocations(), HeapInfoA[HEAP_UNIT_32].MaxAllocations,
			HeapInfoA[HEAP_UNIT_64].CurrentUsage/1024, HeapInfoA[HEAP_UNIT_64].MaxMemory/1024, HeapInfoA[HEAP_UNIT_64].GetAllocations(), HeapInfoA[HEAP_UNIT_64].MaxAllocations,
			HeapInfoA[HEAP_UNIT_256].CurrentUsage/1024, HeapInfoA[HEAP_UNIT_256].MaxMemory/1024, HeapInfoA[HEAP_UNIT_256].GetAllocations(), HeapInfoA[HEAP_UNIT_256].MaxAllocations);
	return Buffer;
}
unsigned int WiiMMGR_AllAllocations(WiiAllocation*	AllocationA)
{
	// this is going to get ugly
	unsigned int Count = 0;
	AllocationA = NULL;
	
#ifdef ENABLE_MEMORY_DEBUG
	WiiMMGR_PrintAll();
	
/*	MemoryTrack* pNode = pBaseTrack;
	while (pNode)
	{
		Count++;
		pNode = pNode->pNext;
		OSOutput(
	}
	// doesn't exist?
	if (Count == 0) return 0;
	
	// allocate nodes
	AllocationA = new WiiAllocation[Count];
	int			i = 0;
	pNode = pBaseTrack;
	while (pNode)
	{
		if (pNode->pFile)
			AllocationA[i].File = pNode->pFile->Text;
		else
			AllocationA[i].File = 0;

		if (pNode->pFunc)
			AllocationA[i].Func = pNode->pFunc->Text;
		else
			AllocationA[i].Func = 0;

		AllocationA[i].Size = pNode->Size;
		AllocationA[i].Line = pNode->Line;

		i++;
		pNode = pNode->pNext;
	}*/
#endif
	return Count;	// total number of nodes?
}
void*	WiiMMGR_ScratchPad()
{
	static void*	pData = MMGR_AllocEx(HEAP_MEM1, 4 * 1024 * 1024, 32);
	return pData;
}


// undefine these if they already exist
#undef	new
#undef	delete

void*	operator new(size_t reportedSize)
{
	return MMGR_AllocEx(HEAP_MEM2,reportedSize, DEFAULT_ALIGNMENT);
}
void*	operator new[](size_t reportedSize)
{
	return MMGR_AllocEx(HEAP_MEM2,reportedSize, DEFAULT_ALIGNMENT);
}
void	operator delete(void *reportedAddress) throw()
{
	MMGR_Free(reportedAddress);
}
void	operator delete[](void *reportedAddress) throw()
{
	MMGR_Free(reportedAddress);
}

// override malloc and free
extern "C" {
void*	malloc(size_t size)
{
	return MMGR_Alloc(size);
}

void free(void *ptr)
{
	MMGR_Free(ptr);
}

void *calloc(size_t nmemb, size_t elemsize)
{
	return MMGR_Alloc(nmemb*elemsize);
}

void *realloc(void *ptr, size_t size)
{
	return MMGR_Realloc(ptr, 0, size, DEFAULT_ALIGNMENT);
}
}

#endif //_PLATFORM_WII
